package utils;

import java.util.BitSet;

import cern.colt.list.ObjectArrayList;



//inspired by http://www.yeminghu.com/?p=124

public class BitSetTrie<T> {

	private BitSetTrieNodeNaive<T> root;
	
	public BitSetTrie(){
		this.root= new BitSetTrieNodeNaive<T>(true,null);
	}
	
	public void put(BitSet bin, T value){
		if(bin!=null){
			root.put(bin, value);
		}
	}
	
	public boolean contains(BitSet bin){
		return root.contains(bin);
	}
	
	public void remove(BitSet bin){
		root.remove(bin);
	}
	
	public T get(BitSet bin){
		return root.get(bin);
	}
	
	public int size(){
		return root.countPrefixes(new BitSet());
	}
	
}

class BitSetTrieNode<T> {
	private BitSet key;
	private boolean marker;
	private T content;
	private int prefixes;
	private BitSetTrieNode<T> parent;
	private ObjectArrayList children;
	
	
	public BitSetTrieNode(BitSet key, BitSetTrieNode<T> parent){
		this.key= key;
		this.marker=false;
		this.content= null;
		this.prefixes= 0;
		this.parent= parent;
		children= new ObjectArrayList();
	}
	
	/**
	 * 
	 * @param bin
	 * @return first mis-matching position, -1 if all match
	 */
	private int matchKey(BitSet bin){
		int a=key.nextSetBit(0),
				b=bin.nextSetBit(0);
		while(a>-1 && b>-1){
			if(a<b){
				return a;
			}
			if(b<a){
				return b;
			}
			a= key.nextSetBit(a+1);
			b= key.nextSetBit(b+1);
		}
		return -1;
	}
	
	public void put(BitSet bin, T value){
		if(bin.isEmpty()){
			marker= true;
			content= value;
			++prefixes;
		}else{
			++prefixes;
			BitSetTrieNode<T> child=null;
			for(int i=0;i<children.size();++i){
				child= (BitSetTrieNode<T>) children.get(i);
				int firstMissPos=child.matchKey(bin);
				if(firstMissPos<0){
					//full length match
					if(bin.length()>=child.key.length()){
						child.put(bin.get(child.key.length(), bin.length()), value);
						break;
					}else{
						//split key, key longer than bin
						BitSetTrieNode<T> inbetween= new BitSetTrieNode<T>(child.key.get(bin.length(), child.key.length()), this);
						child.parent=inbetween;
						inbetween.children.add(child);
						inbetween.prefixes=child.prefixes+1;
						inbetween.marker=true;
						inbetween.content= value;
						children.set(i, inbetween);
						break;
					}
				}else if(firstMissPos>0){
					//split key
					BitSetTrieNode<T> inbetween= new BitSetTrieNode<T>(child.key.get(0, firstMissPos), this);
					child.parent=inbetween;
					child.key=child.key.get(firstMissPos, child.key.length());
					inbetween.children.add(child);
					inbetween.prefixes=child.prefixes;
					inbetween.put(bin.get(firstMissPos, bin.length()), value);
					children.set(i, inbetween);
					break;
				}else{
					child=null;
				}
			}
			if(child==null){
				child= new BitSetTrieNode<T>(bin, this);
				child.put(new BitSet(),value);
				children.add(child);
			}
		}
	}
	
	public boolean contains(BitSet bin){
		if(bin.isEmpty()){
			return marker;
		}else{
			
			
			
			int ch= bin.get(0)?1:0;
			BitSetTrieNode<T> child= children[ch];
			if(child == null){
				return false;
			}else{
				return child.contains(bin.get(1,bin.length()));
			}
		}
	}
	
	public void remove(BitSet bin){
		if(bin.isEmpty()){
			if(marker){
				content= null;
				prefixes--;
			}
			correctTrieNode();
		}else{
			int ch= bin.get(0)?1:0;
			BitSetTrieNode<T> child= children[ch];
			if(child != null) {
				child.remove(bin.get(1, bin.length()));
			}
		}
	}
	
	private void correctTrieNode() {
		if( parent!=null){
			if(content == null && prefixes == 0){
				parent.children[key?1:0] = null;
			}
			--parent.prefixes;
			parent.correctTrieNode();
		}
	}
	
	public T get(BitSet bin){
		if(bin.isEmpty()){
			return content;
		}else{
			int ch= bin.get(0)?1:0;
			BitSetTrieNode<T> child= children[ch];
			if(child == null){
				return null;
			}else{
				return child.get(bin.get(1, bin.length()));
			}
		}
	}
	
	public int countPrefixes(BitSet bin){
		if(bin.isEmpty()){
			return prefixes;
		}else {
			int ch= bin.get(0)?1:0;
			BitSetTrieNode<T> child= children[ch];
			if(child == null){
				return 0;
			}else{
				return child.countPrefixes(bin.get(1, bin.length()));
			}
		}
	}
	
	public boolean getKey(){
		return key;
	}
	
}


class BitSetTrieNodeNaive<T> {
	private boolean key;
	private boolean marker;
	private T content;
	private int prefixes;
	private BitSetTrieNodeNaive<T> parent;
	private BitSetTrieNodeNaive<T> children[];
	
	
	public BitSetTrieNodeNaive(boolean key, BitSetTrieNodeNaive<T> parent){
		this.key= key;
		this.marker=false;
		this.content= null;
		this.prefixes= 0;
		this.parent= parent;
		children= new BitSetTrieNodeNaive[2];
	}
	
	public void put(BitSet bin, T value){
		if(bin.isEmpty()){
			marker= true;
			content= value;
			++prefixes;
		}else{
			++prefixes;
			int ch= bin.get(0)?1:0;
			BitSetTrieNodeNaive<T> child= children[ch];
			if(child==null){
				child= new BitSetTrieNodeNaive<T>(ch==1, this);
				children[ch]=child;
			}
			child.put(bin.get(1,bin.length()),value);
		}
	}
	
	public boolean contains(BitSet bin){
		if(bin.isEmpty()){
			return marker;
		}else{
			int ch= bin.get(0)?1:0;
			BitSetTrieNodeNaive<T> child= children[ch];
			if(child == null){
				return false;
			}else{
				return child.contains(bin.get(1,bin.length()));
			}
		}
	}
	
	public void remove(BitSet bin){
		if(bin.isEmpty()){
			if(marker){
				content= null;
				prefixes--;
			}
			correctTrieNode();
		}else{
			int ch= bin.get(0)?1:0;
			BitSetTrieNodeNaive<T> child= children[ch];
			if(child != null) {
				child.remove(bin.get(1, bin.length()));
			}
		}
	}
	
	private void correctTrieNode() {
		if( parent!=null){
			if(content == null && prefixes == 0){
				parent.children[key?1:0] = null;
			}
			--parent.prefixes;
			parent.correctTrieNode();
		}
	}
	
	public T get(BitSet bin){
		if(bin.isEmpty()){
			return content;
		}else{
			int ch= bin.get(0)?1:0;
			BitSetTrieNodeNaive<T> child= children[ch];
			if(child == null){
				return null;
			}else{
				return child.get(bin.get(1, bin.length()));
			}
		}
	}
	
	public int countPrefixes(BitSet bin){
		if(bin.isEmpty()){
			return prefixes;
		}else {
			int ch= bin.get(0)?1:0;
			BitSetTrieNodeNaive<T> child= children[ch];
			if(child == null){
				return 0;
			}else{
				return child.countPrefixes(bin.get(1, bin.length()));
			}
		}
	}
	
	public boolean getKey(){
		return key;
	}
	
}
